import pandas as pd
import sqlite3
import re
import jieba
from typing import List, Dict
import logging
import os

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class DataCleaner:
    def __init__(self, db_path: str = '../weibo_comments.db'):
        if not os.path.exists(db_path):
            raise FileNotFoundError(f"数据库文件 {db_path} 不存在")
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path)
        logger.info(f"成功连接到数据库 {db_path}")
        
    def check_tables(self) -> List[str]:
        """检查数据库中的所有表"""
        cursor = self.conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = [table[0] for table in cursor.fetchall()]
        logger.info(f"数据库中的表: {tables}")
        return tables
        
    def load_data_from_db(self) -> pd.DataFrame:
        """从数据库加载评论数据"""
        try:
            tables = self.check_tables()
            if 'comments' not in tables:
                raise ValueError("数据库中没有找到comments表")
                
            query = "SELECT * FROM comments"
            logger.info(f"执行查询: {query}")
            df = pd.read_sql_query(query, self.conn)
            logger.info(f"成功从数据库表 comments 加载 {len(df)} 条评论数据")
            return df
        except Exception as e:
            logger.error(f"从数据库加载数据失败: {str(e)}")
            raise

    def clean_text(self, text: str) -> str:
        """清理文本数据"""
        if not isinstance(text, str):
            return ""
            
        # 移除URL
        text = re.sub(r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+', '', text)
        
        # 移除特殊字符和表情符号
        text = re.sub(r'[^\u4e00-\u9fa5a-zA-Z0-9]', ' ', text)
        
        # 移除多余空格
        text = re.sub(r'\s+', ' ', text).strip()
        
        return text

    def process_data(self, df: pd.DataFrame) -> pd.DataFrame:
        """处理数据框中的所有文本"""
        df['cleaned_text'] = df['content'].apply(self.clean_text)
        df['word_count'] = df['cleaned_text'].apply(lambda x: len(x))
        
        # 移除空文本
        df = df[df['cleaned_text'].str.len() > 0]
        
        logger.info(f"数据清洗完成，剩余 {len(df)} 条有效评论")
        return df

    def save_to_csv(self, df: pd.DataFrame, output_path: str = 'cleaned_comments.csv'):
        """保存清洗后的数据到CSV文件"""
        try:
            # 直接保存文件，不需要创建目录
            df.to_csv(output_path, index=False, encoding='utf-8')
            logger.info(f"数据已保存到 {output_path}")
        except Exception as e:
            logger.error(f"保存CSV文件失败: {str(e)}")
            raise

    def save_to_db(self, df: pd.DataFrame, table_name: str = 'cleaned_comments'):
        """保存清洗后的数据到数据库"""
        try:
            df.to_sql(table_name, self.conn, if_exists='replace', index=False)
            logger.info(f"数据已保存到数据库表 {table_name}")
        except Exception as e:
            logger.error(f"保存到数据库失败: {str(e)}")
            raise

    def close(self):
        """关闭数据库连接"""
        self.conn.close()

def main():
    cleaner = DataCleaner()
    try:
        # 加载数据
        df = cleaner.load_data_from_db()
        
        # 清洗数据
        cleaned_df = cleaner.process_data(df)
        
        # 保存到CSV
        cleaner.save_to_csv(cleaned_df)
        
        # 保存到数据库
        cleaner.save_to_db(cleaned_df)
        
    except Exception as e:
        logger.error(f"处理过程中发生错误: {str(e)}")
    finally:
        cleaner.close()

if __name__ == "__main__":
    main() 